import os
import cv2
import numpy as np
import SimpleITK as stk
from time import time
from enum import Enum,unique
from sklearn.model_selection import train_test_split
from keras.utils import to_categorical
from .function import utils

@unique
class Task(Enum):
    Lung_GTV = 3
    Thoracic_OAR = 4

class data:
    def __init__(self,source_path,task_num,width,height):
        self.src = source_path
        self.flag = task_num
        self.width = width
        self.height = height

    def dataset2D(self,width,height):
        train_path  = os.path.join(self.src,'train',Task(self.flag).name)
        list_sorted = sorted(os.listdir(train_path),key=lambda x: int(x))
        self.data_list = []
        self.label_list= []
        for name in list_sorted:
            data    = utils().Read_Image(os.path.join(train_path,name,'data.nii.gz'))
            label   = utils().Read_Image(os.path.join(train_path,name,'label.nii.gz'))
            # clip
            data_clip  = data[np.ix_(range(0,data.shape[0]),range(112, 432),range(90, 410))]
            label_clip = label[np.ix_(range(0,label.shape[0]),range(112, 432),range(90, 410))]
            # 不换行输出方式
            # print('\r',"{}-{}".format(data.shape,label.shape),end='',flush=True)
            for index in range(data_clip.shape[0]):
                slice_data = cv2.resize(data_clip[index],dsize=(width,height))
                slice_label= cv2.resize(label_clip[index],dsize=(width,height))
                self.data_list.append(slice_data)
                self.label_list.append(slice_label)

        self.data_np = np.array(self.data_list)
        self.label_np= np.array(self.label_list)
        return self.data_np,self.label_np

    def trainTest(self,rate,n_classes):
        '''
        Rate is the rate of the test in the dataset.
        '''
        self.dataset2D(self.width,self.height)
        self.data_np  = np.expand_dims(self.data_np,axis=-1)
        self.label_np = utils().one_hot(self.label_np,n_classes)
        return train_test_split(self.data_np,self.label_np,test_size=rate)

if __name__ == "__main__":
    dataset = data('../dataset',3)
    a = time()
    # print(dataset.dataset2D(3)[0].shape)
    # print(len(dataset.trainTest(0.2)))
    b = time()
    print(rate)
    print("{}".format(b-a))
    # print(dir(Task(3)))
    # print(Task(3).value)